#include <cassert>
#include "llvm/ADT/ScopedHashTable.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
#include "mlir/ExecutionEngine/ExecutionEngine.h"
#include "mlir/ExecutionEngine/OptUtils.h"
#include "mlir/IR/Operation.h"
#include "mlir/Target/LLVMIR/ModuleTranslation.h"
#include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
#include "mlir/Dialect/LLVMIR/Transforms/Passes.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir-c/IR.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir-c/Pass.h"
#include "mlir/IR/DialectRegistry.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/Dialect/Func/Extensions/AllExtensions.h"
#include "mlir/Dialect/Affine/Passes.h"
#include "Ch6/include/toy/Dialect.h"
#include "Ch6/include/toy/Passes.h"
#include "cpp_side.h"

void mlirLoadToyDialect(MlirContext rawContext) {
  auto context = (mlir::MLIRContext *)rawContext.ptr;
  assert(context != nullptr && "context is null");

  context->getOrLoadDialect<mlir::toy::ToyDialect>();
}

void *mlirScopedHashTableCreate() {
  return new llvm::ScopedHashTable<llvm::StringRef, mlir::Value>();
}

void mlirScopedHashTableDestroy(void *scopedHashTable) {
  assert(scopedHashTable != nullptr && "scopedHashTable is null");

  auto sht =
      (llvm::ScopedHashTable<llvm::StringRef, mlir::Value> *)scopedHashTable;
  assert(sht != nullptr && "convert scopedHashTable error");

  delete sht;
}

const void *mlirScopedHashTableScopeCreate(void *scopedHashTable) {
  assert(scopedHashTable != nullptr && "scopedHashTable is null");

  auto sht =
      (llvm::ScopedHashTable<llvm::StringRef, mlir::Value> *)scopedHashTable;
  assert(sht != nullptr && "convert scopedHashTable error");

  return new llvm::ScopedHashTableScope<llvm::StringRef, mlir::Value>(*sht);
}

void mlirScopedHashTableScopeDestroy(const void *scopedHashTableScope) {
  assert(scopedHashTableScope != nullptr && "scopedHashTableScope is null");

  auto scope =
      (const llvm::ScopedHashTableScope<llvm::StringRef, mlir::Value> *)
          scopedHashTableScope;
  assert(scope != nullptr);

  delete scope;
}

int mlirDeclare(void *scopedHashTable, const char *name, size_t length, MlirValue rawValue) {
  assert(scopedHashTable != nullptr && "scopedHashTable is null");
  assert(name != nullptr && "name is null");
  assert(!mlirValueIsNull(rawValue) && "rawValue is null");

  auto sht =
      (llvm::ScopedHashTable<llvm::StringRef, mlir::Value> *)scopedHashTable;
  assert(sht != nullptr && "convert scopedHashTable error");

  auto key = llvm::StringRef(name, length);

  if (sht->count(key)) {
    return 0;
  }

  auto value = mlir::Value::getFromOpaquePointer(rawValue.ptr);

  sht->insert(key, value);
  return 1;
}

MlirValue mlirLookup(void *scopedHashTable, const char *name, size_t length) {
  assert(scopedHashTable != nullptr && "scopedHashTable is null");
  assert(name != nullptr && "name is null");

  auto sht =
      (llvm::ScopedHashTable<llvm::StringRef, mlir::Value> *)scopedHashTable;
  assert(sht != nullptr && "convert scopedHashTable error");

  auto key = llvm::StringRef(name, length);
  auto value = sht->lookup(key);
  return MlirValue{value.getAsOpaquePointer()};
}

void *mlirDynCastShapeInference(MlirOperation rawOp) {
  auto op = (mlir::Operation *)rawOp.ptr;
  assert(op != nullptr && "convert Operation error");

  auto shape_op = llvm::dyn_cast<mlir::toy::ShapeInference>(op);
  return shape_op;
}

void mlirShapeInferenceInferShapes(void *shapeInference) {
  assert(shapeInference != nullptr && "shapeInference is null");

  auto op = (mlir::Operation *)shapeInference;
  assert(op != nullptr && "convert Operation error");

  auto shape_op = llvm::dyn_cast<mlir::toy::ShapeInference>(op);
  assert(shape_op != nullptr && "convert ShapeInference error");

  shape_op.inferShapes();
}

void mlirSetPrivate(MlirOperation rawOp) {
  auto op = (mlir::Operation *)rawOp.ptr;
  if (auto funcOp = llvm::dyn_cast<mlir::toy::FuncOp>(op)) {
    funcOp.setPrivate();
  }
}

void mlirFuncRegisterAllExtensions(MlirDialectRegistry rawRegistry) {
  auto registry = (mlir::DialectRegistry *)rawRegistry.ptr;
  assert(registry != nullptr && "convert DialectRegistry error");

  mlir::func::registerAllExtensions(*registry);
}

MlirPass mlirToyCreateLowerToAffinePass() {
  auto pass = mlir::toy::createLowerToAffinePass();
  auto pass_ptr = pass.release();
  return MlirPass{pass_ptr};
}

MlirPass mlirAffineCreateLoopFusionPass() {
  auto pass = mlir::affine::createLoopFusionPass();
  auto pass_ptr = pass.release();
  return MlirPass{pass_ptr};
}

MlirPass mlirAffineCreateAffineScalarReplacementPass() {
  auto pass = mlir::affine::createAffineScalarReplacementPass();
  auto pass_ptr = pass.release();
  return MlirPass{pass_ptr};
}

int mlirDumpLLVMIR(MlirOperation rawOperation, int enableOpt) {
  auto operation = (mlir::Operation *)rawOperation.ptr;
  
  llvm::LLVMContext llvmContext;
  auto llvmModule = mlir::translateModuleToLLVMIR(operation, llvmContext);
  if (!llvmModule) {
    llvm::errs() << "Failed to emit LLVM IR\n";
    return -1;
  }

  // Initialize LLVM targets.
  llvm::InitializeNativeTarget();
  llvm::InitializeNativeTargetAsmPrinter();

  // Configure the LLVM Module
  auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
  if (!tmBuilderOrError) {
    llvm::errs() << "Could not create JITTargetMachineBuilder\n";
    return -1;
  }

  auto tmOrError = tmBuilderOrError->createTargetMachine();
  if (!tmOrError) {
    llvm::errs() << "Could not create TargetMachine\n";
    return -1;
  }
  mlir::ExecutionEngine::setupTargetTripleAndDataLayout(llvmModule.get(),
                                                        tmOrError.get().get());

  /// Optionally run an optimization pipeline over the llvm module.
  auto optPipeline = mlir::makeOptimizingTransformer(
      /*optLevel=*/enableOpt ? 3 : 0, /*sizeLevel=*/0,
      /*targetMachine=*/nullptr);
  if (auto err = optPipeline(llvmModule.get())) {
    llvm::errs() << "Failed to optimize LLVM IR " << err << "\n";
    return -1;
  }
  llvm::errs() << *llvmModule << "\n";
  return 0;
}

int mlirRunJit(MlirOperation rawOperation, int enableOpt) {
  auto module = (mlir::Operation *)rawOperation.ptr;

  // Initialize LLVM targets.
  llvm::InitializeNativeTarget();
  llvm::InitializeNativeTargetAsmPrinter();

  // Register the translation from MLIR to LLVM IR, which must happen before we
  // can JIT-compile.
  mlir::registerBuiltinDialectTranslation(*module->getContext());
  mlir::registerLLVMDialectTranslation(*module->getContext());

  // An optimization pipeline to use within the execution engine.
  auto optPipeline = mlir::makeOptimizingTransformer(
      /*optLevel=*/enableOpt ? 3 : 0, /*sizeLevel=*/0,
      /*targetMachine=*/nullptr);

  // Create an MLIR execution engine. The execution engine eagerly JIT-compiles
  // the module.
  mlir::ExecutionEngineOptions engineOptions;
  engineOptions.transformer = optPipeline;
  auto maybeEngine = mlir::ExecutionEngine::create(module, engineOptions);
  assert(maybeEngine && "failed to construct an execution engine");
  auto &engine = maybeEngine.get();

  // Invoke the JIT-compiled function.
  auto invocationResult = engine->invokePacked("main");
  if (invocationResult) {
    llvm::errs() << "JIT invocation failed\n";
    return -1;
  }

  return 0;
}

void mlirLLVMRegisterInlinerInterface(MlirDialectRegistry rawRegistry) {
  auto registry = (mlir::DialectRegistry *)rawRegistry.ptr;

  mlir::LLVM::registerInlinerInterface(*registry);
}

MlirPass mlirToyCreateLowerToLLVMPass() {
  auto pass = mlir::toy::createLowerToLLVMPass();
  auto pass_ptr = pass.release();
  return MlirPass{pass_ptr};
}

MlirPass mlirLLVMCreateDIScopeForLLVMFuncOpPass() {
  auto pass = mlir::LLVM::createDIScopeForLLVMFuncOpPass();
  auto pass_ptr = pass.release();
  return MlirPass{pass_ptr};
}
